import numpy as np
import matplotlib.mlab as mlab
import matplotlib.pyplot as plt
import math
#计时器
from timeit import default_timer as timer

class GMM:
    def __init__(self,gaussian_num = 10):
        self.gaussian_num = gaussian_num
        self.weights = np.random.uniform(0,1,(gaussian_num))
        self.weights = self.weights / np.sum(self.weights)
        
        self.mu = np.random.uniform(-4,4,(gaussian_num))
        self.sigma = np.random.uniform(0.5,2,(gaussian_num))

    #正态分布
    def normal(self,mu,sigma,weight,x):
        return weight/(np.sqrt(2 * np.pi) * sigma) * np.exp(-(x - mu) * (x - mu) / ( 2 * sigma * sigma))

    #采样
    def sample(self,number = 1000):
        sample = np.array(0)
        for i in range(len(self.weights)):
            weight = math.ceil(number * self.weights[i])
            sample = np.append(sample, np.random.normal(self.mu[i],self.sigma[i],weight))

        #洗牌
        np.random.shuffle(sample)
        return sample
    
    #展示
    def show(self,model = None):
        
        s = [self.normal(self.mu, self.sigma, self.weights, x * 0.01) for x in range(-1000,1000)]

        s = np.sum(s,axis = 1)

        plt.hist(self.sampleData, bins=1000, normed=True)

        x_list = [x * 0.01 for x in range(-1000,1000)]
        plt.plot(x_list, s ,color='green' )

        if model != None:
            s = [self.normal(model.mu, model.sigma, model.weights, x * 0.01) for x in range(-1000,1000)]

            s = np.sum(s,axis = 1)

            x_list = [x * 0.01 for x in range(-1000,1000)]
            plt.plot(x_list, s ,color='red' )

        plt.show()

    def train(self,sample,time):
        
        self.sampleData = sample

        start = timer()
        for index in range(time):
            N = len(sample)

            _gamma = [self.normal(self.mu, self.sigma, self.weights, y)  for y in sample]
            _gamma = [item / np.sum(item) for item in  _gamma]

            _sum_gamma = np.sum(_gamma,axis=0)
            temp = [ _gamma[i] * sample[i] for i in range(N) ]
            _mu = np.sum(temp ,axis=0) / _sum_gamma

            temp = [ _gamma[i] * (sample[i] - self.mu) * (sample[i] - self.mu) for i in range(len(sample)) ]
            _sigma = np.sum(temp ,axis=0)/_sum_gamma

            _weight = _sum_gamma/N

            self.mu =_mu
            self.sigma = np.clip(_sigma,0.1,3)
            self.weights = _weight
        run_time = timer() - start
        print("cpu run time %f seconds " % (run_time))



if __name__ == '__main__':
    gmm = GMM(10)
    #gmm.show()
    gmm2 = GMM(20)
    gmm2.train(gmm.sample(1000000),10)
    gmm2.show(gmm)
